!> author: 左志华
!> date: 2022/05/13
!>
!> 卡方检验
module chi2test_m

    use gsl_randist_m
    use, intrinsic :: iso_fortran_env, only: wp => real64
    implicit none
    private

    public :: chi2test, t_test, wp

    !> 卡方检验
    interface chi2test
        module procedure chi2test_1, chi2test_2
    end interface chi2test

    !> T 检验
    interface t_test
        module procedure welch_t_test, student_t_test
    end interface

contains

    !> 卡方检验 2
    impure subroutine chi2test_2(x, pval, chisq, dof)
        real(wp), intent(in) :: x(:, :)
        real(wp), intent(out) :: pval     !! p-value
        real(wp), intent(out) :: chisq    !! 卡方值
        integer, intent(out) :: dof       !! 自由度
        integer :: s(2)

        s = shape(x)
        if (s(1) == 1) then      ! 单列向量
            dof = s(2) - 1
            chisq = chisq_1(x(1, :))
            pval = gsl_cdf_chisq_q(chisq, real(dof, c_double))
            return
        elseif (s(2) == 1) then  ! 单行向量
            dof = s(1) - 1
            chisq = chisq_1(x(:, 1))
            pval = gsl_cdf_chisq_q(chisq, real(dof, c_double))
            return
        end if

        dof = (s(1) - 1)*(s(2) - 1)
        chisq = chisq_2(x)
        pval = gsl_cdf_chisq_q(chisq, real(dof, c_double))

    end subroutine chi2test_2

    !> 卡方检验 1
    impure subroutine chi2test_1(x, pval, chisq, dof)
        real(wp), intent(in) :: x(:)
        real(wp), intent(out) :: pval     !! p-value
        real(wp), intent(out) :: chisq    !! 卡方值
        integer, intent(out) :: dof       !! 自由度

        chisq = chisq_1(x)
        dof = size(x, 1) - 1
        pval = gsl_cdf_chisq_q(chisq, real(dof, c_double))

    end subroutine chi2test_1

    !> 计算卡方值 2
    pure real(wp) function chisq_2(x) result(chi2)
        real(wp), intent(in) :: x(:, :)
        real(wp) :: expected(size(x, 1), size(x, 2))
        integer :: i, j
        real(wp) :: tol, xtol(size(x, 2)), ytol(size(x, 1))

        xtol = sum(x, 1)
        ytol = sum(x, 2)
        tol = sum(xtol)

        do i = 1, size(x, 1)
            do j = 1, size(x, 2)
                expected(i, j) = xtol(j)*ytol(i)/tol
            end do
        end do

        chi2 = sum((x - expected)**2/expected)

    end function chisq_2

    !> 计算卡方值 1
    pure real(wp) function chisq_1(x) result(chisq)
        real(wp), intent(in) :: x(:)
        real(wp) :: expected

        expected = sum(x)/size(x)
        chisq = sum((x - expected)**2)/expected

    end function chisq_1

    !> Welch T 检验
    impure subroutine welch_t_test(x, y, pval, t, dof)
        real(wp), intent(in) :: x(:), y(:)
        real(wp), intent(out) :: pval       !! p-value
        real(wp), intent(out) :: t          !! T 值
        real(wp), intent(out) :: dof        !! 自由度
        integer :: n1, n2
        real(wp) :: m1, m2, v1, v2

        n1 = size(x)
        n2 = size(y)
        m1 = sum(x)/n1
        m2 = sum(y)/n2
        v1 = sum((x - m1)**2)/(n1 - 1)
        v2 = sum((y - m2)**2)/(n2 - 1)

        t = (m1 - m2)/sqrt(v1/n1 + v2/n2)
        dof = (v1/n1 + v2/n2)**2/(v1**2/(n1**2*(n1 - 1)) + v2**2/(n2**2*(n2 - 1)))
        pval = 2*gsl_cdf_tdist_p(-abs(t), dof)

    end subroutine welch_t_test

    !> 单样本 Student T 检验
    impure subroutine student_t_test(x, mean, pval, t, dof)
        real(wp), intent(in) :: x(:)
        real(wp), intent(in) :: mean        !! 均值
        real(wp), intent(out) :: pval       !! p-value
        real(wp), intent(out) :: t          !! T 值
        integer, intent(out) :: dof         !! 自由度
        integer :: n
        real(wp) :: m, se

        n = size(x)
        m = sum(x)/n
        se = sqrt(sum((x - m)**2)/(n - 1))/sqrt(real(n, wp))
        t = (m - mean)/se                   ! T = (样本均值 - 总体均值) / 标准误差
        dof = n - 1
        pval = 2*gsl_cdf_tdist_q(abs(t), real(dof, c_double))

    end subroutine student_t_test

end module chi2test_m
